'''对下载好的图片进行预处理工作：
1.图像转化为灰度图
2.图像转化为二值图
3.图像降噪，降噪之后再次转化为二值图
4.提取图像轮廓
'''
import  cv2
import  numpy as np
import os
import pathlib
import copy
from my_util import divide_two_box
from my_util import divide_three_box


def image_processing(filepath):
    '''
    图像预处理
    :param filepath 图像存储路径
    :return: contours   轮廓坐标
                : im    图像
                :im_res 二值化处理后的图像
    '''
    im=cv2.imread(filepath)
    im_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    ret, im_inv = cv2.threshold(im_gray,127,255,cv2.THRESH_BINARY_INV)  #黑白二值反转
    kernel = 1/16*np.array([[1,2,1], [2,4,2], [1,2,1]])
    im_blur = cv2.filter2D(im_inv,-1,kernel)    #高斯模糊进行降噪
    ret, im_res = cv2.threshold(im_blur,127,255,cv2.THRESH_BINARY)  #降噪之后再次进行二值处理
    _, contours, hierarchy = cv2.findContours(im_res, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) #找出外轮廓
    return contours,im,im_res

def save_processed_image(filepath,filename):
    '''
    将要预测的验证码处理后的灰度图、二值图放入文件夹
    :param filepath:
    :return:
    '''
    path=r"{}\{}".format(filepath[:len(filepath)-4],filename)
    print(path)
    if not os.path.exists(path):
        os.makedirs(path)
    im=cv2.imread(filepath)
    im_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    gray_path=r"{}\gray.jpg".format(path)
    cv2.imwrite(gray_path,im_gray)

    ret, im_inv = cv2.threshold(im_gray, 127, 255, cv2.THRESH_BINARY_INV)  # 黑白二值反转
    binary_path = r"{}\binary.jpg".format(path)
    #print(binary_path)
    cv2.imwrite(binary_path, im_inv)

    kernel = 1 / 16 * np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]])
    im_blur = cv2.filter2D(im_inv, -1, kernel)  # 高斯模糊进行降噪
    ret, im_res = cv2.threshold(im_blur, 127, 255, cv2.THRESH_BINARY)  # 降噪之后再次进行二值处理
    denoise_path = r"{}\denoise.jpg".format(path)
    cv2.imwrite(denoise_path, im_res)

def image_segmentation(image_dire_path,image_segmented_dire):
    '''
    分割图像
    :param image_path: 存放图像文件夹位置
    :param image_segmented_path: 图像分割后存放文件夹名称（相对路径）
    :return:
    '''
    data_root = pathlib.Path(image_dire_path)
    for item in data_root.iterdir():
        image_rel = pathlib.Path(item).relative_to(data_root)  # 计算相对路径
        contours, im, im_res = image_processing(str(item))
        rect_points = []  # 存放切割出来的矩形的四个点坐标
        contour_length = len(contours)
        w_lis = []  # 存放矩形轮廓的宽度
        ################################
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            w_lis.append(w)
        #print("contours_length", len(contours))
        for contour in contours:

            if contour_length == 6:
                x, y, w, h = cv2.boundingRect(contour)
                rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                rect_points.extend(rect_point)

            if contour_length == 5:
                x, y, w, h = cv2.boundingRect(contour)
                if w == max(w_lis):  # w_max是所有contonur的宽度中最宽的值
                    rect_left, rect_right = divide_two_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_right)
                else:
                    x, y, w, h = cv2.boundingRect(contour)
                    rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                    rect_points.extend(rect_point)

            if contour_length == 4:
                x, y, w, h = cv2.boundingRect(contour)
                if w == max(w_lis) and w > 2 * min(w_lis):  # w_max是所有contonur的宽度中最宽的值
                    rect_left, rect_mid, rect_right = divide_three_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_mid)
                    rect_points.extend(rect_right)
                elif w <= 2 * min(w_lis) and w > min(w_lis):
                    rect_left, rect_right = divide_two_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_right)
                else:
                    x, y, w, h = cv2.boundingRect(contour)
                    rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                    rect_points.extend(rect_point)

            if contour_length == 3:
                x, y, w, h = cv2.boundingRect(contour)
                if w > 2 * min(w_lis) and 2 * w > 5 * min(w_lis):
                    rect_left, rect_mid, rect_right = divide_three_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_mid)
                    rect_points.extend(rect_right)
                elif w > min(w_lis) and w <= 2 * min(w_lis):
                    rect_left, rect_right = divide_two_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_right)
                else:
                    x, y, w, h = cv2.boundingRect(contour)
                    rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                    rect_points.extend(rect_point)

            if contour_length == 2:
                x, y, w, h = cv2.boundingRect(contour)
                rect_left, rect_mid, rect_right = divide_three_box(x, y, w, h)
                rect_points.extend(rect_left)
                rect_points.extend(rect_mid)
                rect_points.extend(rect_right)

        # print("rect_points_length",len(rect_points))
        if (len(rect_points) == 24):
            '对轮廓部分进行排序，让识别处理的图片按照从左往右的顺序排列'
            x_points = np.array(copy.deepcopy(rect_points))
            x_points = x_points[np.argsort(x_points[:, 0])]
            y_points = np.array(copy.deepcopy(rect_points))
            y_points = y_points[np.argsort(y_points[:, 1])]
            for i in range(6):
                img_segmented = im_res[y_points[0][1]:y_points[23][1], x_points[i * 4][0]:x_points[i * 4 + 3][0]]
                # img_segmented = im_res[rect_points[i * 4][1]:rect_points[i * 4 + 3][1],rect_points[i * 4][0]:rect_points[i * 4 + 1][0]]
                if img_segmented != []:
                    img_segmented = cv2.resize(img_segmented, (25, 25))  # 将字符图片统一调整为25x25的图片大小
                    filename = "{}.jpg".format(str(image_rel)[:6] + "-" + str(i) + "-" + str(image_rel)[i])
                    print(filename)
                    filepath = os.path.join(image_segmented_dire, filename)
                    cv2.imwrite(filepath, img_segmented)

def train_image_segmentation():
    '''
    训练集图像分割
    :return:
    '''
    image_segmentation(r"C:\Users\14121\Desktop\train","train_segmented_image")


def test_image_segmentation():
    '''
    测试图像分割
    :return:
    '''
    image_segmentation(r"C:\Users\14121\Desktop\test","test_segmented_image")



def predict_image_segmentation():
    '''
        预测图像分割
        :return:
        '''
    data_root = pathlib.Path(r"C:\Users\14121\Desktop\predict")
    for item in data_root.iterdir():
        print(item)
        image_rel = pathlib.Path(item).relative_to(data_root)  # 计算相对路径
        contours, im, im_res = image_processing(str(item))
        rect_points = []  # 存放切割出来的矩形的四个点坐标
        contour_length = len(contours)
        #print("未处理前",contour_length)
        w_lis = []  # 存放矩形轮廓的宽度
        for contour in contours:
                x, y, w, h = cv2.boundingRect(contour)
                #cv2.rectangle(bgr_img, (x, y), (x + w, y + h), (0, 255, 0), 1)
                w_lis.append(w)
        #cv2.imwrite('contours.png', bgr_img)

        for contour in contours:
            if contour_length == 6 :
                x, y, w, h = cv2.boundingRect(contour)
                rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                rect_points.extend(rect_point)

            if contour_length == 5:
                x, y, w, h = cv2.boundingRect(contour)
                if w == max(w_lis):  # w_max是所有contonur的宽度中最宽的值
                    rect_left, rect_right = divide_two_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_right)
                else:
                    x, y, w, h = cv2.boundingRect(contour)
                    rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                    rect_points.extend(rect_point)

            if contour_length == 4:
                x, y, w, h = cv2.boundingRect(contour)
                if w == max(w_lis) and w > 2 * min(w_lis):  # w_max是所有contonur的宽度中最宽的值
                    rect_left, rect_mid, rect_right = divide_three_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_mid)
                    rect_points.extend(rect_right)
                elif w <= 2 * min(w_lis) and w > min(w_lis):
                    rect_left, rect_right = divide_two_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_right)
                else:
                    x, y, w, h = cv2.boundingRect(contour)
                    rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                    rect_points.extend(rect_point)

            if contour_length == 3:
                x, y, w, h = cv2.boundingRect(contour)
                if w > 2 * min(w_lis) and 2 * w > 5 * min(w_lis):
                    rect_left, rect_mid, rect_right = divide_three_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_mid)
                    rect_points.extend(rect_right)
                elif w > min(w_lis) and w <= 2 * min(w_lis):
                    rect_left, rect_right = divide_two_box(x, y, w, h)
                    rect_points.extend(rect_left)
                    rect_points.extend(rect_right)
                else:
                    x, y, w, h = cv2.boundingRect(contour)
                    rect_point = np.int64([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
                    rect_points.extend(rect_point)

            if contour_length == 2:
                x, y, w, h = cv2.boundingRect(contour)
                rect_left, rect_mid, rect_right = divide_three_box(x, y, w, h)
                rect_points.extend(rect_left)
                rect_points.extend(rect_mid)
                rect_points.extend(rect_right)
        print("rect_point", rect_points)
        if (len(rect_points) == 24):
            '对轮廓部分进行排序，让识别处理的图片按照从左往右的顺序排列'
            x_points = np.array(copy.deepcopy(rect_points))
            x_points = x_points[np.argsort(x_points[:, 0])]
            y_points = np.array(copy.deepcopy(rect_points))
            y_points = y_points[np.argsort(y_points[:, 1])]
            for i in range(6):
                img_segmented = im_res[y_points[0][1]:y_points[23][1],
                                x_points[i * 4][0]:x_points[i * 4 + 3][0]]
                img_segmented = cv2.resize(img_segmented, (25, 25))  # 将字符图片统一调整为25x25的图片大小
                filename = "{}.jpg".format(str(image_rel)[:6] + "-" + str(i))
                print(filename)
                filepath = r"D:\dataset\predict_segmented_image\{}".format(filename)
                cv2.imwrite(filepath, img_segmented)

def get_train_img():
    '''
    获取训练的图片，并将其放到数组里面
    :return:训练图片列表
    '''
    img_lis = []    #存放图片
    data_root = pathlib.Path(r"D:\dataset\train_segmented_image")
    for item in data_root.iterdir():
        im = cv2.imread(str(item))
        img_lis.append(im)
    return np.array(img_lis)/255.0  #归一化

def get_predict_img():
    '''
    获取预测分割后的图片，并将其放到数组里面
    :return:训练图片列表
    '''
    img_lis = []    #存放图片
    data_root = pathlib.Path(r"D:\dataset\predict_segmented_image")
    for item in data_root.iterdir():
        im = cv2.imread(str(item))
        img_lis.append(im)
    return np.array(img_lis)/255.0  #归一化

def get_train_label():
    '''
    获取训练图片的标签，并将其放到数组里面
    :return: 训练图片标签
    '''
    img_label_lis=[]    #存放图片的label
    data_root=pathlib.Path(r"D:\dataset\train_segmented_image")
    for item in data_root.iterdir():
        image_rel = pathlib.Path(item).relative_to(data_root)  # 计算相对路径
        image_rel=str(image_rel)
        image_rel_length=len(image_rel)
        img_label=image_rel[image_rel_length-5]
        img_label_lis.append(img_label)
    return img_label_lis

def get_test_img():
    '''
    获取测试的图片，并将其放到数组里面
    :return:测试图片列表
    '''
    img_lis = []    #存放图片
    data_root = pathlib.Path(r"D:\dataset\test_image")
    for item in data_root.iterdir():
        im = cv2.imread(str(item))
        img_lis.append(im)
    return np.array(img_lis)/255.0  #归一化

def get_test_label():
    '''
    获取测试图片的标签，并将其放到数组里面
    :return:测试图片标签
    '''
    img_label_lis=[]    #存放图片的label
    data_root=pathlib.Path(r"D:\dataset\test_label")
    for item in data_root.iterdir():
        image_rel = pathlib.Path(item).relative_to(data_root)  # 计算相对路径
        image_rel=str(image_rel)
        image_rel_length=len(image_rel)
        img_label=image_rel[image_rel_length-5]
        img_label_lis.append(img_label)
    return img_label_lis



if __name__=='__main__':

    pass
    #test_image_segmentation()
    #train_image_segmentation()
    #print(get_test_label())
    #print(get_train_img().shape)
    #print(get_train_label())
    #save_processed_image(r"C:\Users\14121\Pictures\python.jpg","test")